import sys
sys.path.append("./")
import time
import cv2
import numpy as np
import functools
import torch

from PyQt5.QtWidgets import *
from PyQt5.QtCore import Qt

from PyQt5.QtGui import QImage, QPixmap, QPainter


class QPixmapLabel(QLabel):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.pixmap_original = None  # 原始图片
        self.scale_factor = 1.0  # 初始缩放比例

    def setPixmap(self, pixmap):
        """ 重新设置原始 Pixmap，并重置缩放比例 """
        self.pixmap_original = pixmap
        self.scale_factor = 1.0  # 重置缩放比例
        self.update_pixmap()

    def update_pixmap(self):
        """ 根据 scale_factor 更新 QLabel 显示的图片 """
        if self.pixmap_original:
            scaled_pixmap = self.pixmap_original.scaled(
                self.pixmap_original.size() * self.scale_factor,
                Qt.KeepAspectRatio,
                Qt.SmoothTransformation
            )
            super().setPixmap(scaled_pixmap)

    def wheelEvent(self, event):
        """ 处理鼠标滚轮缩放 """
        if event.angleDelta().y() > 0:  # 滚轮向上，放大
            self.scale_factor *= 1.1
        else:  # 滚轮向下，缩小
            self.scale_factor /= 1.1

        self.scale_factor = max(0.2, min(self.scale_factor, 5.0))  # 限制缩放范围
        self.update_pixmap()



class CircuitBreakerApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.initUI()
        self.initFlags()
        self.initModel()
        self.camera_opened = False

    def initUI(self):
        self.setWindowTitle("低压断路器状态检测系统")
        self.setGeometry(100, 100, 1280, 800)

        # 主布局（左右结构）
        main_widget = QWidget()
        main_layout = QHBoxLayout()
        main_widget.setLayout(main_layout)
        self.setCentralWidget(main_widget)

        # ===== 左侧功能按钮区域 =====
        left_panel = QWidget()
        left_panel.setFixedWidth(200)
        btn_layout = QVBoxLayout()
        btn_layout.setSpacing(15)
        
        # 功能按钮
        buttons = [
            ("图片检测", self.load_image),
            ("视频检测", self.load_video),
            ("打开摄像头", self.toggle_camera),
            ("摄像头检测", self.start_detect),
            ("关闭摄像头", self.close_camera)
        ]
        
        for text, handler in buttons:
            btn = QPushButton(text)
            btn.setFixedHeight(40)
            btn.clicked.connect(handler)
            btn_layout.addWidget(btn)
        
        left_panel.setLayout(btn_layout)
        main_layout.addWidget(left_panel)

        # ===== 右侧显示区域 =====
        right_panel = QWidget()
        right_layout = QVBoxLayout()
        right_panel.setLayout(right_layout)

        
        # 1. 图像显示区域（占50%高度）
        # 在显示区域下方添加三个按钮
        button_layout = QHBoxLayout()

        self.btn3 = QPushButton("运行")

        # 设定按钮大小
        self.btn3.setFixedSize(80, 40)

        # 设定按钮的点击事件（可根据需要修改）
        self.btn3.clicked.connect(self.on_button3_click)

        # 将按钮添加到水平布局中，并调整位置
        button_layout.addWidget(self.btn3, alignment=Qt.AlignRight)
        # 将按钮布局添加到右侧主布局
        right_layout.addLayout(button_layout)

        display_title = QLabel("显示区域")
        display_title.setStyleSheet("font-weight: bold; color: rgb(0, 0, 0);")
        # self.img_display = QLabel("请先加载图片或视频") # 初始文本加大字体
        self.img_display = QPixmapLabel()
        self.img_display.setAlignment(Qt.AlignCenter)
        # 白色背景，黑色字体
        self.img_display.setStyleSheet("""
            background-color: rgb(255, 255, 255);
            color: rgb(0, 0, 0);
            border: 2px solid #3E3E3E;
            border-radius: 30px;
            """)
        self.img_display.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        right_layout.addWidget(display_title)
        right_layout.addWidget(self.img_display, stretch=8)

        # 2. 检测结果区域（占30%高度）
        result_group = QWidget()
        result_layout = QVBoxLayout()
        result_group.setLayout(result_layout)
        
        result_title = QLabel("检测结果：")
        result_title.setStyleSheet("font-weight: bold; color: rgb(0, 0, 0);")
        self.result_area = QTextEdit()
        self.result_area.setReadOnly(True)
        self.result_area.setStyleSheet("""
            background-color: rgb(255, 255, 255);
            color: rgb(0, 0, 0);
            border: 1px solid #3E3E3E;
            padding: 10px;
        """)
        
        result_layout.addWidget(result_title)
        result_layout.addWidget(self.result_area)
        right_layout.addWidget(result_group, stretch=1)


        # 3. 打印输出区域（占20%高度）
        output_group = QWidget()
        output_layout = QVBoxLayout()
        output_group.setLayout(output_layout)
        
        output_title = QLabel("系统输出：")
        output_title.setStyleSheet("font-weight: bold; color: rgb(0, 0, 0);")
        self.output_area = QTextEdit()
        self.output_area.setReadOnly(True)
        self.output_area.setPlainText("特别说明：如需启动检测，请先加载model文件！！！")
        self.output_area.setStyleSheet("""
            background-color: rgb(255, 255, 255);
            color: rgb(0, 0, 0);
            border: 1px solid #3E3E3E;
            padding: 10px;
        """)
        
        output_layout.addWidget(output_title)
        output_layout.addWidget(self.output_area)
        right_layout.addWidget(output_group, stretch=1)

        main_layout.addWidget(right_panel)

    def initFlags(self):
        self.image = None
        self.image_path = None

    def initModel(self):
        from infer_model import create_model, run_model

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = create_model(device=device)
        self.infer_func = functools.partial(run_model, model, device)

        current_time = time.strftime("%Y-%m-%d, %H:%M:%S", time.localtime())
        self.output_area.append(f"{current_time}\t [系统] 模型加载成功！")

    # ===== 功能函数 =====
    # 定义按钮的槽函数
    def on_button3_click(self):
        current_time = time.strftime("%Y-%m-%d, %H:%M:%S", time.localtime())
        if self.image is None:
            self.output_area.append(f"{current_time}\t [系统] 请先加载图片或视频！")
            return

        self.output_area.append(f"{current_time}\t [系统] 开始检测: {self.image_path}")

        # inference
        self.result_area.clear()
        t_start = time.time()
        result = self.infer_func(self.image)
        t_end = time.time()
        result_text = f"""
            <span style="font-size:15px;">检测数据:\t\t </span> 
            <span style="color:blue; font-size:15px;">{self.image_path}</span><br>

            <span style="font-size:15px;">检测结果:\t\t </span> 
            <span style="color:{'green' if result == 'ON' else 'red'}; font-size:16px; font-weight:bold;">
                {result}
            </span><br>

            <span style="font-size:15px;">检测耗时:\t\t </span> 
            <span style="color:purple; font-size:15px;">{round((t_end - t_start)*1000, 1)} ms</span>
        """
        self.result_area.setHtml(result_text)



        self.output_area.append(f"{current_time}\t [系统] 检测结束")

    # 摄像头相关函数
    def toggle_camera(self):
        # 先清空历史显示
        self.result_area.clear()

        self.camera_opened = not self.camera_opened
        status = "开启" if self.camera_opened else "关闭"
        current_time = time.strftime("%Y-%m-%d, %H:%M:%S", time.localtime())
        self.output_area.append(f"{current_time}\t [系统] 摄像头状态：{status}")
        
        if self.camera_opened:
            self.show_placeholder("实时摄像头画面")
        else:
            self.img_display.clear()
        

    # 关闭摄像头
    def close_camera(self):
        if self.camera_opened:
            self.toggle_camera()

    # 图片和视频加载
    def load_image(self):
        # 先清空历史显示
        self.result_area.clear()
        # self.output_area.clear()

        # 弹出文件对话框，选择图片
        options = QFileDialog.Options()
        options |= QFileDialog.ReadOnly
        file_path, _ = QFileDialog.getOpenFileName(self, "选择图片", "", "Images (*.png *.jpg *.jpeg *.bmp)", options=options)
        image = cv2.imdecode(np.fromfile(file_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        self.image = image
        self.image_path = file_path

        # 显示图片
        if image is not None:
            height, width, channel = image.shape
            bytes_per_line = 3 * width
            q_image = QImage(image.data, width, height, bytes_per_line, QImage.Format_RGB888).rgbSwapped()
            pixmap = QPixmap.fromImage(q_image)

            max_size = 416
            # 如果宽度或高度大于最大值，进行缩放
            if pixmap.width() > max_size or pixmap.height() > max_size:
                pixmap = pixmap.scaled(max_size, max_size, Qt.KeepAspectRatio)

            self.img_display.setPixmap(pixmap)
            self.img_display.setAlignment(Qt.AlignCenter)
        else:
            self.img_display.setText("加载失败")
            

        current_time = time.strftime("%Y-%m-%d, %H:%M:%S", time.localtime())
        self.output_area.append(f"{current_time}\t [系统] 已加载图片: {file_path}")

    # 加载视频
    def load_video(self):
        # 先清空历史显示
        self.result_area.clear()
        # self.output_area.clear()

        self.show_placeholder("视频播放")
        current_time = time.strftime("%Y-%m-%d, %H:%M:%S", time.localtime())
        self.output_area.append(f"{current_time}\t [系统] 已加载示例视频")


    # 开始检测
    def start_detect(self):
        # 先清空历史显示
        self.result_area.clear()
        # self.output_area.clear()

        current_time = time.strftime("%Y-%m-%d, %H:%M:%S", time.localtime())
        self.output_area.append(f"{current_time}\t [系统] 开始检测...")
        # 模拟检测结果
        self.result_area.append("设备状态：正常\n温度：45°C\n电流：32A")


    # 显示占位图像
    def show_placeholder(self, text):
        """显示占位图像"""
        pixmap = QPixmap(800, 500)
        pixmap.fill(Qt.darkGray)
        painter = QPainter(pixmap)
        painter.setPen(Qt.white)
        painter.setFont(painter.font())
        painter.drawText(pixmap.rect(), Qt.AlignCenter, text)
        self.img_display.setPixmap(pixmap)
        painter.end()

if __name__ == "__main__":
    app = QApplication(sys.argv)
    app.setStyle("Fusion")  # 使用现代样式
    window = CircuitBreakerApp()
    window.show()
    sys.exit(app.exec_())